#include "stdafx.h"
#include "DriverWrapper.h"
#include "Message.h"
#include "ProcessMgr.h"
#include "EventMgr.h"
#include "../Shared/Protocol.h"
#include "util.h"
#include "nt.h"

CMessageParsers m_MessageParser;

class CMessageParser_PsCreateProcess : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_ps_create_process_data *data = (svc_ps_create_process_data *)buf;

		FILETIME createTime = *(FILETIME *)&data->CreateTime;

		if (data->Create)
		{
			m_ProcessMgr.Lock();
			up = m_ProcessMgr.Find(data->ProcessId);
			if (!up)
			{
				up = new CUniqueProcess(data->ProcessId, data->ParentProcessId, &createTime, data->Is64Bit ? true : false, data->SessionId, data->ImagePath);
				m_ProcessMgr.InsertProcess(up);
				m_ProcessMgr.FillIcon(up);
				m_ProcessMgr.FillParent(up);
				m_ProcessMgr.CallWndCallbacks(up, WM_ADDPROCESS);
			}
			m_ProcessMgr.Unlock();
		}
		else
		{
			m_ProcessMgr.Lock();
			up = m_ProcessMgr.Find(data->ProcessId);
			if (up)
			{
				up->m_bAlive = false;
				m_ProcessMgr.CallWndCallbacks(up, WM_KILLPROCESS);
			}
			m_ProcessMgr.Unlock();
		}

		if (up)
		{
			if(data->Create)
				ev = new CUniqueEvent_ProcessCreate(up, data->time);
			else
				ev = new CUniqueEvent_ProcessDestroy(up, data->time);
		}
		return ev;
	}
};

class CMessageParser_PsLoadImage : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_ps_load_image_data *data = (svc_ps_load_image_data *)buf;

		std::wstring ImagePath = data->ImagePath;

		if (ImagePath.empty() && data->ProcessId > 4)//Not a system process...
		{
			TCHAR szImagePath[MAX_PATH];
			if (GetImagePathByAddress(data->ProcessId, data->ImageBase, szImagePath, sizeof(szImagePath)))
				ImagePath = szImagePath;
		}

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			ev = new CUniqueEvent_LoadImage(up, ImagePath.c_str(), data->ImageBase, data->ImageSize, data->time);
		}
		return ev;
	}
};

class CMessageParser_NtLoadDriver : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_load_driver_data *data = (svc_nt_load_driver_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			ev = new CUniqueEvent_LoadDriver(up, data->RegisterPath, data->ImagePath, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_NtQuerySystemInfo : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_query_systeminfo_data *data = (svc_nt_query_systeminfo_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			switch (data->QueryClass)
			{
			case SystemProcessInformation:
				ev = new CUniqueEvent_EnumProcess(up, data->time);
				break;
			case SystemModuleInformation:
				ev = new CUniqueEvent_EnumSystemModule(up, data->time);
				break;
			case SystemHandleInformation:
				ev = new CUniqueEvent_EnumSystemHandle(up, data->time);
				break;
			case SystemObjectInformation:
				ev = new CUniqueEvent_EnumSystemObject(up, data->time);
				break;
			}
		}
		return ev;
	}
};

class CMessageParser_NtOpenProcess : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueProcess *target = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_open_process_data *data = (svc_nt_open_process_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		target = m_ProcessMgr.Find(data->TargetProcessId);
		m_ProcessMgr.Unlock();

		if (up && target)
		{
			ev = new CUniqueEvent_OpenProcess(up, target, data->DesiredAccess, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_NtTerminateProcess : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueProcess *target = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_terminate_process_data *data = (svc_nt_terminate_process_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		target = m_ProcessMgr.Find(data->TargetProcessId);
		m_ProcessMgr.Unlock();

		if (up && target)
		{
			ev = new CUniqueEvent_TerminateProcess(up, target, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_NtReadVirtualMemory : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueProcess *target = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_read_virtual_mem_data *data = (svc_nt_read_virtual_mem_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		target = m_ProcessMgr.Find(data->TargetProcessId);
		m_ProcessMgr.Unlock();

		if (up && target)
		{
			ev = new CUniqueEvent_ReadVirtualMemory(up, target, data->BaseAddress, data->BufferSize, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_NtWriteVirtualMemory : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueProcess *target = NULL;
		CUniqueEvent *ev = NULL;
		svc_nt_write_virtual_mem_data *data = (svc_nt_write_virtual_mem_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		target = m_ProcessMgr.Find(data->TargetProcessId);
		m_ProcessMgr.Unlock();

		if (up && target)
		{
			ev = new CUniqueEvent_WriteVirtualMemory(up, target, data->BaseAddress, data->BufferSize, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_FsCreateFile : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_fs_create_file_data *data = (svc_fs_create_file_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			ev = new CUniqueEvent_CreateFile(up, data->FilePath, data->DesiredAccess, data->Disposition, data->Options, data->ShareAccess, data->Attributes, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_FsCloseFile : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_fs_close_file_data *data = (svc_fs_close_file_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			ev = new CUniqueEvent_CloseFile(up, data->FilePath, data->ResultStatus, data->time);
		}
		return ev;
	}
};

class CMessageParser_FsReadWriteFile : public CMessageParser
{
public:
	CUniqueEvent *ParseBytes(PUCHAR buf)
	{
		CUniqueProcess *up = NULL;
		CUniqueEvent *ev = NULL;
		svc_fs_readwrite_file_data *data = (svc_fs_readwrite_file_data *)buf;

		m_ProcessMgr.Lock();
		up = m_ProcessMgr.Find(data->ProcessId);
		m_ProcessMgr.Unlock();

		if (up)
		{
			if(data->IsWrite)
				ev = new CUniqueEvent_WriteFile(up, data->FilePath, data->ReadLength, data->ByteOffset, data->ResultStatus, data->time);
			else
				ev = new CUniqueEvent_ReadFile(up, data->FilePath, data->ReadLength, data->ByteOffset, data->ResultStatus, data->time);
		}
		return ev;
	}
};

void InitMessageParser(void)
{
	m_MessageParser[svc_ps_create_process] = new CMessageParser_PsCreateProcess;
	m_MessageParser[svc_ps_load_image] = new CMessageParser_PsLoadImage;
	m_MessageParser[svc_nt_load_driver] = new CMessageParser_NtLoadDriver;
	m_MessageParser[svc_nt_query_systeminfo] = new CMessageParser_NtQuerySystemInfo;
	m_MessageParser[svc_nt_open_process] = new CMessageParser_NtOpenProcess;
	m_MessageParser[svc_nt_terminate_process] = new CMessageParser_NtTerminateProcess;
	m_MessageParser[svc_nt_read_virtual_mem] = new CMessageParser_NtReadVirtualMemory;
	m_MessageParser[svc_nt_write_virtual_mem] = new CMessageParser_NtWriteVirtualMemory;
	m_MessageParser[svc_fs_create_file] = new CMessageParser_FsCreateFile;
	m_MessageParser[svc_fs_close_file] = new CMessageParser_FsCloseFile;
	m_MessageParser[svc_fs_readwrite_file] = new CMessageParser_FsReadWriteFile;
}

bool ParseMessage(PUCHAR data)
{
	int protocol = data[0];

	CMessageParsers::iterator itor = m_MessageParser.find(protocol);
	if (itor != m_MessageParser.end())
	{
		CUniqueEvent *ev = itor->second->ParseBytes(data);
		if (ev)
		{
			m_EventMgr.Lock();
			m_EventMgr.InsertEvent(ev);
			m_EventMgr.Unlock();
			m_EventMgr.CallWndCallbacks(ev, WM_ADDEVENT);			
			return true;
		}
	}

	return false;
}